# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openFuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
This module unit-tests NPUCollector's DCMI-backed helpers for device temperature,
health state, and aggregated error codes.
"""

import unittest
import ctypes
from ctypes import c_int, c_uint
from unittest.mock import patch, Mock

from hardware_monitor.collector.metrics.npu.collector_for_npu import NPUCollector
from hardware_monitor.collector.metrics.npu.constants import (RET_NO_FUNCTION, RET_INVALID_VALUE, INVALID_TEMPERATURE,
                                                              MAX_ERROR_CODE_COUNT, DEVICE_NOT_READY_ERROR_CODE,
                                                              CARD_DROP_FAULT_CODE)


class TestGetDeviceTemperature(unittest.TestCase):
    """
    TestGetDeviceTemperature tests `_get_device_temperature` for symbol presence,
    error returns, value checks, and success.
    """

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_func_missing(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_temperature behavior
        when the DcmiManager instance is missing.
        """
        mock_get_dcmi_instance.return_value = None

        # pylint: disable=protected-access
        res = NPUCollector._get_device_temperature(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], RET_NO_FUNCTION)
        self.assertIn("not available", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_ret_error(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_temperature behavior
        when the Dcmi function returns an error code.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        mock_dcmi_instance.dcmi.dcmi_get_device_temperature.return_value = -8001

        # pylint: disable=protected-access
        res = NPUCollector._get_device_temperature(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], -8001)
        self.assertIn("Failed to get temperature for logicID(0)", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_temp_invalid(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_temperature behavior
        when the temperature value is invalid.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake(card, dev, temp_ref):  # pylint: disable=unused-argument
            ctypes.cast(temp_ref, ctypes.POINTER(c_int)).contents.value = INVALID_TEMPERATURE - 1
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_temperature.side_effect = fake

        # pylint: disable=protected-access
        res = NPUCollector._get_device_temperature(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], RET_INVALID_VALUE)
        self.assertIn("Invalid temperature value for logicID(0)", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_success(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_temperature behavior
        when the temperature value is successfully retrieved.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake(card, dev, temp_ref):  # pylint: disable=unused-argument
            ctypes.cast(temp_ref, ctypes.POINTER(c_int)).contents.value = 67
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_temperature.side_effect = fake

        # pylint: disable=protected-access
        res = NPUCollector._get_device_temperature(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res, {"temperature": 67})


class TestGetDeviceHealth(unittest.TestCase):
    """
    TestGetDeviceHealth tests `_get_device_health` for symbol presence, error handling,
    flag validation, and success mapping
    """

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_func_missing(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_health behavior
        when the DcmiManager instance is missing.
        """
        mock_get_dcmi_instance.return_value = None

        # pylint: disable=protected-access
        res = NPUCollector._get_device_health(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], RET_NO_FUNCTION)
        self.assertIn("not available", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_ret_error(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_health behavior
        when the Dcmi function returns an error code.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        mock_dcmi_instance.dcmi.dcmi_get_device_health.return_value = -8001

        # pylint: disable=protected-access
        res = NPUCollector._get_device_health(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], -8001)
        self.assertIn("ret: -8001", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_health_invalid_flag(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_health behavior
        when an invalid health flag is returned.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_ok(card, dev, health_ref):  # pylint: disable=unused-argument
            ctypes.cast(health_ref, ctypes.POINTER(c_uint)).contents.value = 2 ** 31 - 1
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_ok

        # pylint: disable=protected-access
        res = NPUCollector._get_device_health(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], RET_INVALID_VALUE)
        self.assertIn("get wrong health state", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_health_success_zero(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_health behavior
        when the health flag is 0 (indicating healthy device).
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_ok(card, dev, health_ref):  # pylint: disable=unused-argument
            ctypes.cast(health_ref, ctypes.POINTER(c_uint)).contents.value = 0
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_ok

        # pylint: disable=protected-access
        res = NPUCollector._get_device_health(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res, {"health": 0})

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_health_success_nonzero(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_health behavior
        when the health flag is non-zero (indicating faulty device).
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_ok(card, dev, health_ref):  # pylint: disable=unused-argument
            ctypes.cast(health_ref, ctypes.POINTER(c_uint)).contents.value = 3
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_ok

        # pylint: disable=protected-access
        res = NPUCollector._get_device_health(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res, {"health": 1})


class TestGetDeviceAllErrorCode(unittest.TestCase):
    """
    TestGetDeviceAllErrorCode tests `_get_device_all_error_code` for error-paths, health interplay,
    count validation, and success.
    """

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_func_missing(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_all_error_code behavior
        when the DcmiManager instance is missing.
        """
        mock_get_dcmi_instance.return_value = None

        # pylint: disable=protected-access
        res = NPUCollector._get_device_all_error_code(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], RET_NO_FUNCTION)
        self.assertIn("not available", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_ret_error(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_all_error_code behavior
        when the Dcmi function returns an error code.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_errcodes(card, dev, err_cnt_ref, code_arr, max_cnt):  # pylint: disable=unused-argument
            return -8001

        def fake_health(card, dev, health_ref):  # pylint: disable=unused-argument
            ctypes.cast(health_ref, ctypes.POINTER(c_uint)).contents.value = 0
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_errorcode_v2.side_effect = fake_errcodes
        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_health

        # pylint: disable=protected-access
        res = NPUCollector._get_device_all_error_code(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], -8001)
        self.assertIn("Failed to obtain the device error code", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_device_not_ready_card_drop(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_all_error_code behavior
        when the device health indicates a card drop fault.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_errcodes(card, dev, err_cnt_ref, code_arr, max_cnt):  # pylint: disable=unused-argument
            ctypes.cast(err_cnt_ref, ctypes.POINTER(c_int)).contents.value = 0
            return 0

        def fake_health(card, dev, health_ref):  # pylint: disable=unused-argument
            return DEVICE_NOT_READY_ERROR_CODE

        mock_dcmi_instance.dcmi.dcmi_get_device_errorcode_v2.side_effect = fake_errcodes
        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_health

        # pylint: disable=protected-access
        res = NPUCollector._get_device_all_error_code(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_count"], 1)
        self.assertEqual(res["error_codes"], [int(CARD_DROP_FAULT_CODE)])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_error_count_invalid(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_all_error_code behavior
        when the error count exceeds the maximum allowed.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_errcodes(card, dev, err_cnt_ref, code_arr, max_cnt):  # pylint: disable=unused-argument
            ctypes.cast(err_cnt_ref, ctypes.POINTER(c_int)).contents.value = MAX_ERROR_CODE_COUNT + 1
            return 0

        def fake_health(card, dev, health_ref):  # pylint: disable=unused-argument
            ctypes.cast(health_ref, ctypes.POINTER(c_uint)).contents.value = 0
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_errorcode_v2.side_effect = fake_errcodes
        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_health

        # pylint: disable=protected-access
        res = NPUCollector._get_device_all_error_code(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_code"], RET_INVALID_VALUE)
        self.assertIn("invalid error code count", res["error_message"])

    @patch("hardware_monitor.collector.metrics.npu.dcmi.dcmi.DcmiManager.get_dcmi_instance")
    def test_success(self, mock_get_dcmi_instance):
        """
        Tests NPUCollector._get_device_all_error_code behavior
        when device error codes are successfully retrieved.
        """
        mock_dcmi_instance = Mock()
        mock_get_dcmi_instance.return_value = mock_dcmi_instance

        def fake_errcodes(card, dev, err_cnt_ref, code_arr, max_cnt):  # pylint: disable=unused-argument
            ctypes.cast(err_cnt_ref, ctypes.POINTER(c_int)).contents.value = 3
            code_arr[0] = 100
            code_arr[1] = 200
            code_arr[2] = 0
            return 0

        def fake_health(card, dev, health_ref):  # pylint: disable=unused-argument
            ctypes.cast(health_ref, ctypes.POINTER(c_uint)).contents.value = 0
            return 0

        mock_dcmi_instance.dcmi.dcmi_get_device_errorcode_v2.side_effect = fake_errcodes
        mock_dcmi_instance.dcmi.dcmi_get_device_health.side_effect = fake_health

        # pylint: disable=protected-access
        res = NPUCollector._get_device_all_error_code(logic_id=0, card_id=0, device_id=0)
        self.assertEqual(res["error_count"], 3)
        self.assertEqual(res["error_codes"], [100, 200])


if __name__ == "__main__":
    unittest.main()
